package main

import (
	"bufio"
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"golang.org/x/term"
)

var (
	blkidCmd  string
	mountCmd  string
	umountCmd string
	fsckCmd   string
)

const (
	mountBaseDir    = "/mnt/garf-data"
	configFile      = "/mnt/garf-data/data.conf"
	dirPermissions  = 0755
	filePermissions = 0644
)

var systemCriticalMounts = []string{"/boot", "/etc", "/tmp"}

type DiskMapping struct {
	OriginalID string `json:"original_id"`
	ShortName  string `json:"short_name"`
}

func checkCommands() bool {
	commands := map[string]*string{
		"blkid":  &blkidCmd,
		"mount":  &mountCmd,
		"umount": &umountCmd,
		"fsck":   &fsckCmd,
	}

	alternativeCommands := map[string][]string{
		"blkid":  {"lsblk", "fdisk -l"},
		"mount":  {"mount"},
		"umount": {"umount"},
		"fsck":   {"e2fsck", "xfs_repair"},
	}

	for cmd, cmdVar := range commands {
		path, err := exec.LookPath(cmd)
		if err == nil {
			*cmdVar = path
			continue
		}

		for _, altCmd := range alternativeCommands[cmd] {
			path, err := exec.LookPath(altCmd)
			if err == nil {
				*cmdVar = path
				log.Printf("Using %s as an alternative for %s\n", altCmd, cmd)
				break
			}
		}

		if *cmdVar == "" {
			log.Printf("Neither %s nor its alternatives are available\n", cmd)
			return false
		}
	}

	return true
}

func init() {
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)
	if !checkCommands() {
		log.Println("Required commands are not available. Exiting.")
		os.Exit(1)
	}
}

func main() {
	if len(os.Args) > 1 && os.Args[1] == "unmount" {
		if len(os.Args) < 3 {
			log.Println("Usage: mounr unmount <mountpoint>")
			os.Exit(1)
		}
		if err := unmountDisk(os.Args[2]); err != nil {
			log.Fatalf("Failed to unmount: %v", err)
		}
	} else if len(os.Args) > 1 && os.Args[1] == "conflicts" {
		displayMountStatus()
		if err := showConflicts(); err != nil {
			log.Fatalf("Failed to show conflicts: %v", err)
		}
	} else {
		conflictMappings, err := autoMount()
		if err != nil {
			log.Fatalf("Auto mount failed: %v", err)
		}
		displayMountStatus()
		if len(conflictMappings) > 0 {
			log.Println("以下磁盘存在冲突：")
			for _, mapping := range conflictMappings {
				log.Printf("  %s -> %s", mapping.OriginalID, filepath.Join(mountBaseDir, mapping.ShortName))
			}
			log.Println("虽然是冲突，但不一定要处理，亦可能是之前已挂载磁盘的挂载目录或之前已链接磁盘目录的软链，请自行甄别...")
		}
	}
}

func autoMount() ([]DiskMapping, error) {
	disks := getAvailableDisks()
	if len(disks) == 0 {
		return nil, fmt.Errorf("无可用磁盘")
	}

	mappings, err := readConfig()
	if err != nil {
		log.Printf("读取配置文件失败: %v. 将创建新的配置。", err)
		mappings = []DiskMapping{}
	}

	var newMappings []DiskMapping
	var conflictMappings []DiskMapping
	for _, disk := range disks {
		// 排除 data.conf
		if strings.HasSuffix(disk, "data.conf") {
			continue
		}

		shortName := generateShortName(disk)
		mountPoint := filepath.Join(mountBaseDir, shortName)

		existingMountPoint, err := getExistingMountPoint(disk)
		if err == nil {
			// 磁盘已经挂载
			if isCriticalMount(existingMountPoint) {
				log.Printf("跳过系统关键挂载点 %s", existingMountPoint)
				continue
			}
			err := createSymlink(existingMountPoint, mountPoint)
			if err != nil {
				log.Printf("为已挂载磁盘 %s 创建软链接失败: %v", disk, err)
				conflictMappings = append(conflictMappings, DiskMapping{OriginalID: disk, ShortName: shortName})
				continue
			}
			log.Printf("为已挂载磁盘 %s 创建软链接 %s -> %s", disk, mountPoint, existingMountPoint)
			newMappings = append(newMappings, DiskMapping{OriginalID: disk, ShortName: shortName})
			continue
		}

		// 磁盘未挂载，执行挂载操作
		fsType, err := getFileSystemType(disk)
		if err != nil {
			log.Printf("获取文件系统类型失败，跳过 %s: %v", disk, err)
			continue
		}

		if fsType == "swap" {
			log.Printf("跳过交换分区 %s", disk)
			continue
		}

		if err := checkFileSystem(disk, fsType); err != nil {
			log.Printf("文件系统检查失败，跳过 %s: %v", disk, err)
			continue
		}

		if err := os.MkdirAll(mountPoint, dirPermissions); err != nil {
			log.Printf("创建挂载点失败 %s: %v", mountPoint, err)
			continue
		}

		if err := mountDisk(disk, mountPoint); err != nil {
			log.Printf("挂载失败 %s: %v", disk, err)
			continue
		}

		newMappings = append(newMappings, DiskMapping{OriginalID: disk, ShortName: shortName})
		log.Printf("成功挂载 %s 到 %s", disk, mountPoint)
	}

	// 更新配置文件，添加所有新的映射
	updatedMappings := updateMappings(mappings, newMappings)
	if err := saveConfig(updatedMappings); err != nil {
		log.Printf("保存配置文件失败: %v", err)
	}

	return conflictMappings, nil
}

func getAvailableDisks() []string {
	var cmd *exec.Cmd
	switch {
	case strings.Contains(blkidCmd, "blkid"):
		cmd = exec.Command(blkidCmd, "-o", "device")
	case strings.Contains(blkidCmd, "lsblk"):
		cmd = exec.Command(blkidCmd, "-nlo", "NAME")
	case strings.Contains(blkidCmd, "fdisk"):
		cmd = exec.Command("sh", "-c", blkidCmd+" | grep '^/dev/' | cut -d' ' -f1")
	default:
		log.Println("No suitable command found for listing disks")
		return nil
	}

	output, err := cmd.Output()
	if err != nil {
		log.Printf("Failed to get disk list: %v\n", err)
		return nil
	}

	disks := strings.Split(string(output), "\n")
	var availableDisks []string
	for _, disk := range disks {
		if disk != "" && !strings.Contains(disk, "loop") && !strings.Contains(disk, "ram") {
			// 检查磁盘是否已经挂载
			_, err := getExistingMountPoint(strings.TrimSpace(disk))
			if err != nil {
				// 磁盘未挂载，添加到可用磁盘列表
				availableDisks = append(availableDisks, strings.TrimSpace(disk))
			}
		}
	}
	return availableDisks
}

func getExistingMountPoint(disk string) (string, error) {
	content, err := os.ReadFile("/proc/mounts")
	if err != nil {
		return "", fmt.Errorf("读取 /proc/mounts 失败: %v", err)
	}

	lines := strings.Split(string(content), "\n")
	for _, line := range lines {
		fields := strings.Fields(line)
		if len(fields) >= 2 && fields[0] == disk {
			return fields[1], nil
		}
	}

	return "", fmt.Errorf("未找到磁盘 %s 的挂载点", disk)
}

func isCriticalMount(mountPoint string) bool {
	for _, criticalMount := range systemCriticalMounts {
		if mountPoint == criticalMount || strings.HasPrefix(mountPoint, criticalMount+"/") {
			return true
		}
	}
	return false
}

func createSymlink(target, link string) error {
	if isCriticalMount(target) {
		return fmt.Errorf("不能为系统关键目录 %s 创建软链接", target)
	}

	existingTarget, err := os.Readlink(link)
	if err == nil {
		// 软链接已存在
		if existingTarget == target {
			// 软链接已经指向正确的目标
			return nil
		}
		// 软链接指向错误的目标，记录但不修改
		return fmt.Errorf("软链接 %s 已存在但指向 %s 而不是 %s", link, existingTarget, target)
	}

	if os.IsNotExist(err) {
		// 软链接不存在，创建新的
		return os.Symlink(target, link)
	}

	// 路径存在但不是软链接
	return fmt.Errorf("路径 %s 已存在且不是软链接", link)
}

func getFileSystemType(disk string) (string, error) {
	cmd := exec.Command(blkidCmd, "-o", "value", "-s", "TYPE", disk)
	output, err := cmd.Output()
	if err != nil {
		return "", fmt.Errorf("failed to get filesystem type: %v", err)
	}
	return string(bytes.TrimSpace(output)), nil
}

func checkFileSystem(disk, fsType string) error {
	var cmd *exec.Cmd
	switch fsType {
	case "ext2", "ext3", "ext4":
		cmd = exec.Command("e2fsck", "-n", disk) // 只读检查
	case "xfs":
		cmd = exec.Command("xfs_repair", "-n", disk) // 只读检查
	case "btrfs":
		cmd = exec.Command("btrfs", "check", disk)
	case "ntfs":
		cmd = exec.Command("ntfsfix", "-n", disk) // 只读检查
	default:
		return fmt.Errorf("不支持的文件系统类型: %s", fsType)
	}

	output, err := cmd.CombinedOutput()
	if err != nil {
		log.Printf("文件系统检查发现问题 %s (类型: %s): %v\n输出: %s", disk, fsType, err, string(output))
		if askForRepair() {
			return repairFileSystem(disk, fsType)
		}
	}

	log.Printf("文件系统检查完成 %s (类型: %s)", disk, fsType)
	return nil
}

func repairFileSystem(disk, fsType string) error {
	var cmd *exec.Cmd
	switch fsType {
	case "ext2", "ext3", "ext4":
		cmd = exec.Command("e2fsck", "-y", disk)
	case "xfs":
		cmd = exec.Command("xfs_repair", disk)
	case "btrfs":
		cmd = exec.Command("btrfs", "check", "--repair", disk)
	case "ntfs":
		cmd = exec.Command("ntfsfix", disk)
	default:
		return fmt.Errorf("不支持的文件系统类型: %s", fsType)
	}

	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("文件系统修复失败 %s (类型: %s): %v\n输出: %s", disk, fsType, err, string(output))
	}

	log.Printf("文件系统修复完成 %s (类型: %s)", disk, fsType)
	return nil
}

func askForRepair() bool {
	fmt.Print("文件系统检查发现问题。是否进行修复？(y/n): ")
	var response string
	fmt.Scanln(&response)
	return strings.ToLower(response) == "y"
}

func generateShortName(disk string) string {
	hash := md5.Sum([]byte(disk))
	return hex.EncodeToString(hash[:])[:8]
}

func mountDisk(disk, mountPoint string) error {
	fsType, err := getFileSystemType(disk)
	if err != nil {
		return fmt.Errorf("获取文件系统类型失败: %v", err)
	}

	cmd := exec.Command(mountCmd, "-t", fsType, disk, mountPoint)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("挂载失败 %s 到 %s: %v\n输出: %s", disk, mountPoint, err, string(output))
	}
	return nil
}

func unmountDisk(mountPoint string) error {
	cmd := exec.Command(umountCmd, mountPoint)
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to unmount %s: %v", mountPoint, err)
	}
	log.Printf("Successfully unmounted %s\n", mountPoint)

	mappings, err := readConfig()
	if err != nil {
		return fmt.Errorf("failed to read config: %v", err)
	}

	var updatedMappings []DiskMapping
	for _, mapping := range mappings {
		if filepath.Join(mountBaseDir, mapping.ShortName) != mountPoint {
			updatedMappings = append(updatedMappings, mapping)
		}
	}
	return saveConfig(updatedMappings)
}

func readConfig() ([]DiskMapping, error) {
	data, err := os.ReadFile(configFile)
	if err != nil {
		if os.IsNotExist(err) {
			return []DiskMapping{}, nil
		}
		return nil, fmt.Errorf("读取配置失败: %v", err)
	}
	var mappings []DiskMapping
	if err := json.Unmarshal(data, &mappings); err != nil {
		return nil, fmt.Errorf("解析配置失败: %v", err)
	}
	return mappings, nil
}

func saveConfig(mappings []DiskMapping) error {
	var validMappings []DiskMapping
	for _, mapping := range mappings {
		fullPath := filepath.Join(mountBaseDir, mapping.ShortName)
		if _, err := os.Lstat(fullPath); err == nil {
			validMappings = append(validMappings, mapping)
		} else {
			log.Printf("跳过不存在的目录映射: %s -> %s", mapping.OriginalID, mapping.ShortName)
		}
	}

	data, err := json.MarshalIndent(validMappings, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal config: %v", err)
	}
	return os.WriteFile(configFile, data, filePermissions)
}

func updateMappings(oldMappings, newMappings []DiskMapping) []DiskMapping {
	mappingMap := make(map[string]DiskMapping)
	for _, mapping := range oldMappings {
		mappingMap[mapping.OriginalID] = mapping
	}
	for _, mapping := range newMappings {
		mappingMap[mapping.OriginalID] = mapping
	}

	var updatedMappings []DiskMapping
	for _, mapping := range mappingMap {
		updatedMappings = append(updatedMappings, mapping)
	}
	return updatedMappings
}

func showConflicts() error {
	mappings, err := readConfig()
	if err != nil {
		return fmt.Errorf("读取配置失败: %v", err)
	}

	conflicts := false
	for _, mapping := range mappings {
		mountPoint := filepath.Join(mountBaseDir, mapping.ShortName)
		target, err := os.Readlink(mountPoint)
		if err != nil {
			if os.IsNotExist(err) {
				log.Printf("软链接不存在: %s", mountPoint)
			} else {
				log.Printf("检查软链接失败 %s: %v", mountPoint, err)
			}
			conflicts = true
			continue
		}

		existingMountPoint, err := getExistingMountPoint(mapping.OriginalID)
		if err != nil {
			log.Printf("获取磁盘 %s 的挂载点失败: %v", mapping.OriginalID, err)
			conflicts = true
			continue
		}

		if target != existingMountPoint {
			log.Printf("冲突: %s\n  配置的目标: %s\n  实际挂载点: %s", mountPoint, target, existingMountPoint)
			conflicts = true
		}
	}

	if !conflicts {
		log.Println("没有发现冲突")
	}

	return nil
}

func readMountInfo() (map[string]string, error) {
	file, err := os.Open("/proc/mounts")
	if err != nil {
		return nil, err
	}
	defer file.Close()

	mountInfo := make(map[string]string)
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())
		if len(fields) >= 2 {
			mountInfo[fields[1]] = fields[0]
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return mountInfo, nil
}

func getActualMountPoint(path string, mountInfo map[string]string) string {
	if device, ok := mountInfo[path]; ok {
		return device
	}
	return "未挂载"
}

func cleanupMappings() error {
	mappings, err := readConfig()
	if err != nil {
		return fmt.Errorf("读取配置失败: %v", err)
	}

	var updatedMappings []DiskMapping
	for _, mapping := range mappings {
		fullPath := filepath.Join(mountBaseDir, mapping.ShortName)
		if _, err := os.Lstat(fullPath); err == nil {
			updatedMappings = append(updatedMappings, mapping)
		} else {
			log.Printf("移除不存在的映射: %s -> %s", mapping.OriginalID, mapping.ShortName)
		}
	}

	if len(updatedMappings) < len(mappings) {
		if err := saveConfig(updatedMappings); err != nil {
			return fmt.Errorf("保存更新后的配置失败: %v", err)
		}
		log.Println("配置文件已更新，移除了不存在的映射")
	} else {
		log.Println("没有需要清理的映射")
	}

	return nil
}

func displayMountStatus() {
	width, _, err := term.GetSize(int(os.Stdout.Fd()))
	if err != nil {
		log.Printf("无法获取终端大小: %v", err)
		width = 80 // 默认宽度
	}

	entries, err := os.ReadDir(mountBaseDir)
	if err != nil {
		log.Printf("无法读取目录 %s: %v", mountBaseDir, err)
		return
	}

	mountInfo, err := readMountInfo()
	if err != nil {
		log.Printf("无法读取挂载信息: %v", err)
		return
	}

	fmt.Printf("\n%s 目录下的项目状态：\n", mountBaseDir)
	printLine(width)

	nameWidth := 12
	typeWidth := 10
	statusWidth := 8
	// 减少目标/挂载点列的宽度
	targetWidth := (width - nameWidth - typeWidth - statusWidth - 6) * 7 / 8

	fmt.Printf("%-*s %-*s %-*s %-*s\n", nameWidth, "名称", typeWidth, "类型", targetWidth, "目标/挂载点", statusWidth, "状态")
	printLine(width)

	var longPaths []string

	for _, entry := range entries {
		if entry.Name() == "data.conf" {
			continue
		}

		fullPath := filepath.Join(mountBaseDir, entry.Name())
		info, err := os.Lstat(fullPath)
		if err != nil {
			log.Printf("无法获取 %s 的信息: %v", fullPath, err)
			continue
		}

		var itemType, target, status string
		if info.Mode()&os.ModeSymlink != 0 {
			itemType = "软链接"
			target, _ = os.Readlink(fullPath)
		} else if info.IsDir() {
			itemType = "挂载点"
			target = getActualMountPoint(fullPath, mountInfo)
		} else {
			continue
		}

		if target == "未挂载" {
			status = "未挂载"
		} else {
			status = "正常"
		}

		displayTarget := target
		if len(target) > targetWidth {
			displayTarget = target[:targetWidth-3] + "..."
			longPaths = append(longPaths, fmt.Sprintf("%s: %s", entry.Name(), target))
		}

		fmt.Printf("%-*s %-*s %-*s %-*s\n", nameWidth, entry.Name(), typeWidth, itemType, targetWidth, displayTarget, statusWidth, status)
	}

	printLine(width)

	// 打印长路径的详细信息
	if len(longPaths) > 0 {
		fmt.Println("\n长路径详细信息:")
		for _, path := range longPaths {
			fmt.Println(path)
		}
	}
}

func printLine(width int) {
	fmt.Println(strings.Repeat("-", width))
}
